import os
from typing import Tuple, Dict, Any
from pathlib import Path
import streamlit as st

# 支持的文件类型和MIME类型映射
SUPPORTED_TYPES = {
    'pdf': ['application/pdf'],
    'docx': [
        'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        'application/zip'  # docx文件本质上是zip格式
    ],
    'md': ['text/plain', 'text/markdown']
}

# 文件大小限制（MB）
MAX_FILE_SIZE_MB = 200

def validate_file(uploaded_file) -> Tuple[bool, str]:
    """验证上传的文件
    
    Args:
        uploaded_file: Streamlit上传的文件对象
        
    Returns:
        Tuple[bool, str]: (是否有效, 错误信息)
    """
    
    try:
        # 检查文件大小
        file_size_mb = len(uploaded_file.getvalue()) / (1024 * 1024)
        if file_size_mb > MAX_FILE_SIZE_MB:
            return False, f"文件大小超过限制 ({file_size_mb:.2f}MB > {MAX_FILE_SIZE_MB}MB)"
        
        # 检查文件扩展名
        file_extension = Path(uploaded_file.name).suffix.lower().lstrip('.')
        if file_extension not in SUPPORTED_TYPES:
            return False, f"不支持的文件格式: {file_extension}"
        
        # 检查文件内容（MIME类型）
        file_content = uploaded_file.getvalue()
        
        # 对于不同文件类型进行特定验证
        if file_extension == 'pdf':
            if not file_content.startswith(b'%PDF'):
                return False, "文件不是有效的PDF格式"
        
        elif file_extension == 'docx':
            # docx文件是zip格式，检查zip头部
            if not (file_content.startswith(b'PK\x03\x04') or file_content.startswith(b'PK\x05\x06')):
                return False, "文件不是有效的DOCX格式"
        
        elif file_extension == 'md':
            # Markdown文件应该是文本格式
            try:
                file_content.decode('utf-8')
            except UnicodeDecodeError:
                try:
                    file_content.decode('gbk')
                except UnicodeDecodeError:
                    return False, "Markdown文件编码格式不支持"
        
        return True, ""
        
    except Exception as e:
        return False, f"文件验证出错: {str(e)}"

def get_file_info(uploaded_file) -> Dict[str, Any]:
    """获取文件详细信息
    
    Args:
        uploaded_file: Streamlit上传的文件对象
        
    Returns:
        Dict[str, Any]: 文件信息字典
    """
    
    try:
        file_content = uploaded_file.getvalue()
        file_size_mb = len(file_content) / (1024 * 1024)
        file_extension = Path(uploaded_file.name).suffix.lower().lstrip('.')
        
        info = {
            'filename': uploaded_file.name,
            'format': file_extension.upper(),
            'size_bytes': len(file_content),
            'size_mb': file_size_mb,
            'mime_type': uploaded_file.type or 'unknown'
        }
        
        # 根据文件类型获取特定信息
        if file_extension == 'pdf':
            info.update(_get_pdf_info(file_content))
        elif file_extension == 'docx':
            info.update(_get_docx_info(file_content))
        elif file_extension == 'md':
            info.update(_get_markdown_info(file_content))
        
        return info
        
    except Exception as e:
        return {
            'filename': uploaded_file.name,
            'format': 'Unknown',
            'size_mb': 0,
            'error': str(e)
        }

def _get_pdf_info(file_content: bytes) -> Dict[str, Any]:
    """获取PDF文件信息"""
    try:
        import fitz  # PyMuPDF
        
        # 从字节流创建PDF文档
        doc = fitz.open(stream=file_content, filetype="pdf")
        
        info = {
            'pages': doc.page_count,
            'title': doc.metadata.get('title', ''),
            'author': doc.metadata.get('author', ''),
            'subject': doc.metadata.get('subject', ''),
            'creator': doc.metadata.get('creator', ''),
            'encrypted': doc.needs_pass,
            'text_extractable': True
        }
        
        # 检查是否可以提取文本
        if doc.page_count > 0:
            first_page = doc[0]
            text = first_page.get_text()
            info['text_extractable'] = len(text.strip()) > 0
        
        doc.close()
        return info
        
    except Exception as e:
        return {
            'pages': 0,
            'error': f"PDF信息提取失败: {str(e)}"
        }

def _get_docx_info(file_content: bytes) -> Dict[str, Any]:
    """获取DOCX文件信息"""
    try:
        from docx import Document
        from io import BytesIO
        
        # 从字节流创建Document对象
        doc = Document(BytesIO(file_content))
        
        # 统计段落和表格
        paragraphs = len([p for p in doc.paragraphs if p.text.strip()])
        tables = len(doc.tables)
        
        # 获取文档属性
        core_props = doc.core_properties
        
        info = {
            'paragraphs': paragraphs,
            'tables': tables,
            'title': core_props.title or '',
            'author': core_props.author or '',
            'subject': core_props.subject or '',
            'created': core_props.created.isoformat() if core_props.created else '',
            'modified': core_props.modified.isoformat() if core_props.modified else ''
        }
        
        return info
        
    except Exception as e:
        return {
            'paragraphs': 0,
            'tables': 0,
            'error': f"DOCX信息提取失败: {str(e)}"
        }

def _get_markdown_info(file_content: bytes) -> Dict[str, Any]:
    """获取Markdown文件信息"""
    try:
        # 尝试不同编码解码
        text = None
        encoding = 'utf-8'
        
        try:
            text = file_content.decode('utf-8')
        except UnicodeDecodeError:
            try:
                text = file_content.decode('gbk')
                encoding = 'gbk'
            except UnicodeDecodeError:
                text = file_content.decode('utf-8', errors='ignore')
                encoding = 'utf-8 (with errors)'
        
        if text:
            lines = text.split('\n')
            
            # 统计标题层级
            headers = {}
            for line in lines:
                line = line.strip()
                if line.startswith('#'):
                    level = len(line) - len(line.lstrip('#'))
                    if level <= 6:  # Markdown最多6级标题
                        headers[f'h{level}'] = headers.get(f'h{level}', 0) + 1
            
            # 统计其他元素
            paragraphs = len([line for line in lines if line.strip() and not line.strip().startswith('#')])
            code_blocks = text.count('```')
            links = text.count('](') + text.count('http')
            
            info = {
                'lines': len(lines),
                'paragraphs': paragraphs,
                'headers': headers,
                'code_blocks': code_blocks // 2,  # 代码块成对出现
                'links': links,
                'encoding': encoding,
                'characters': len(text)
            }
            
            return info
        
        return {'error': '无法解码文件内容'}
        
    except Exception as e:
        return {
            'lines': 0,
            'paragraphs': 0,
            'error': f"Markdown信息提取失败: {str(e)}"
        }

def check_file_health(uploaded_file) -> Dict[str, Any]:
    """检查文件健康状态
    
    Args:
        uploaded_file: Streamlit上传的文件对象
        
    Returns:
        Dict[str, Any]: 健康检查结果
    """
    
    health_report = {
        'overall_status': 'healthy',
        'issues': [],
        'warnings': [],
        'recommendations': []
    }
    
    try:
        # 基本验证
        is_valid, error_msg = validate_file(uploaded_file)
        if not is_valid:
            health_report['overall_status'] = 'error'
            health_report['issues'].append(error_msg)
            return health_report
        
        # 获取文件信息
        file_info = get_file_info(uploaded_file)
        file_extension = Path(uploaded_file.name).suffix.lower().lstrip('.')
        
        # 文件大小检查
        if file_info['size_mb'] > 50:
            health_report['warnings'].append(f"文件较大 ({file_info['size_mb']:.2f}MB)，处理可能需要较长时间")
        
        # 特定格式检查
        if file_extension == 'pdf':
            if file_info.get('encrypted'):
                health_report['issues'].append("PDF文件已加密，无法处理")
                health_report['overall_status'] = 'error'
            elif not file_info.get('text_extractable'):
                health_report['warnings'].append("PDF可能是扫描版，文本提取效果可能不佳")
                health_report['recommendations'].append("建议使用OCR功能或转换为可编辑格式")
        
        elif file_extension == 'docx':
            if file_info.get('tables', 0) > 10:
                health_report['warnings'].append(f"文档包含大量表格 ({file_info['tables']}个)")
                health_report['recommendations'].append("表格内容将转换为文本格式")
        
        elif file_extension == 'md':
            if 'errors' in file_info.get('encoding', ''):
                health_report['warnings'].append("文件编码可能存在问题")
                health_report['recommendations'].append("建议检查文件编码格式")
        
        # 设置整体状态
        if health_report['issues']:
            health_report['overall_status'] = 'error'
        elif health_report['warnings']:
            health_report['overall_status'] = 'warning'
        
        return health_report
        
    except Exception as e:
        health_report['overall_status'] = 'error'
        health_report['issues'].append(f"健康检查失败: {str(e)}")
        return health_report